% Copyright (C) 2014-19 Benjamin Born and Johannes Pfeifer 
%
% This is free software: you can redistribute it and/or modify
% it under the terms of the GNU General Public License as published by
% the Free Software Foundation, either version 3 of the License, or
% (at your option) any later version.
%
% This code is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
% GNU General Public License for more details.
%
% For a copy of the licencse,
% see <http://www.gnu.org/licenses/>.

%% Figure 3 in the paper

%% Setup

% clear; close all; clc
clear
addpath('../Auxiliary_Files')

unconditional_model_dummy = 0; % 1 for upper panel, 0 for middle and lower panel
sign_dummy = 0; % -1: cut, 0: symmetric

%the naming of the following needs to be consistent with the field names in the pos structure
impulse_name = 'g_growth_fe'; % OE shocks

% Load data arrays for regression
load('../data_array_for_regression_stacked_by_variable_CDS_new')

lag_number = 4;
max_irf_horizon = 8;
time_fixed_effects_dummy = 1;
country_fixed_effects_dummy = 1;

add_residual_dummy = 1; % adds residuals from previous step to regression to increase efficiency;

var_dummy = 1; % baseline
switch var_dummy
    case 1
        regressor_var_names={'g_growth_realization','y_growth_realization','debt_growth_real_demeaned_linearly_detrended'};
        dependent_var_names={'g_growth_realization','y_growth_realization','spread_end_quarter','debt_to_GDP_demeaned_linearly_detrended'}; %make sure spread is consistent with indicator if requested!
        dependent_var_plottitles={'Government spending growth','GDP growth','Default premium','Debt/GDP'};
        dependent_var_ylabels={'percentage points','percentage points','basis points','percentage points'};
    otherwise
        error('Case is not defined')
end

indicator_dummy = 1; % baseline
switch indicator_dummy
    case 1 % 1 = Country specific CDF with spread at the end of quarter
        Indicator_position=pos.empirical_cdf_country_group_specific_end_quarter;
        [regressor_var_names]=check_dependent_and_regressors_for_consistency(dependent_var_names,regressor_var_names,'spread_end_quarter',0,0);
        indicator_save_name='ecdf_gspec_endquart';
        Figure_name_indicator='Empirial CDF Group-Specific End Quarter';
        State_name_vector={'No fiscal stress','Maximum stress'};
    otherwise
        error('Case is not defined')
end
regressor_var_names = {};
Figure_name_split='All Countries';
Figure_name_indicator='Empirial CDF Group-Specific End Quarter';

%% Shock Variable

fe_shocks_temp = data_array_for_regression_stacked_by_variable(:,:,pos.(impulse_name));

fe_shocks_negative = zeros(size(fe_shocks_temp));
fe_shocks_negative(fe_shocks_temp<0 | isnan(fe_shocks_temp)) = fe_shocks_temp(fe_shocks_temp<0 | isnan(fe_shocks_temp));
fe_shocks_positive = zeros(size(fe_shocks_temp));
fe_shocks_positive(fe_shocks_temp>0 | isnan(fe_shocks_temp)) = fe_shocks_temp(fe_shocks_temp>0 | isnan(fe_shocks_temp));
if any(any(fe_shocks_negative>0)) || any(any(fe_shocks_positive<0 ))
    error('Sign is wrong')
end

if unconditional_model_dummy == 1
    if sign_dummy==-1
        data_array_for_regression_impulse_only=cat(3,fe_shocks_negative,fe_shocks_positive);
    elseif sign_dummy==1
        data_array_for_regression_impulse_only=cat(3,fe_shocks_positive,fe_shocks_negative);
    elseif sign_dummy==0
        data_array_for_regression_impulse_only=cat(3,fe_shocks_temp);
    else
        error('undefined case for sign_dummy')
    end
elseif unconditional_model_dummy == 0
    % neg_shock*F(z)
    data_array_for_regression_g_stress_neg=fe_shocks_negative.*...
    data_array_for_regression_stacked_by_variable(:,:,Indicator_position);
    % neg_shock*(1-F(z))
    data_array_for_regression_g_nostress_neg=fe_shocks_negative.*...
    (1-data_array_for_regression_stacked_by_variable(:,:,Indicator_position));

    % pos_shock*F(z)
    data_array_for_regression_g_stress_pos=fe_shocks_positive.*...
    data_array_for_regression_stacked_by_variable(:,:,Indicator_position);
    % pos_shock*(1-F(z))
    data_array_for_regression_g_nostress_pos=fe_shocks_positive.*...
    (1-data_array_for_regression_stacked_by_variable(:,:,Indicator_position));

    if sign_dummy==-1
        data_array_for_regression_impulse_only=cat(3,...
            data_array_for_regression_g_stress_neg,data_array_for_regression_g_nostress_neg,...
            data_array_for_regression_g_stress_pos,data_array_for_regression_g_nostress_pos...
            );
    elseif sign_dummy==1
        data_array_for_regression_impulse_only=cat(3,...
            data_array_for_regression_g_stress_pos,data_array_for_regression_g_nostress_pos,...
            data_array_for_regression_g_stress_neg,data_array_for_regression_g_nostress_neg...
            );
    elseif sign_dummy==0
        % shock*F(z)
        data_array_for_regression_g_stress=fe_shocks_temp.*...
            data_array_for_regression_stacked_by_variable(:,:,Indicator_position);
        % shock*(1-F(z))
        data_array_for_regression_g_nostress=fe_shocks_temp.*...
            (1-data_array_for_regression_stacked_by_variable(:,:,Indicator_position));
        
        data_array_for_regression_impulse_only=cat(3,...
            data_array_for_regression_g_stress,data_array_for_regression_g_nostress...
            );
    end
else
    error('case not defined')
end
data_array_for_regression=data_array_for_regression_impulse_only;

%% Lagged regressors

% ATTENTION: Spreads always last regressor

for var_iter=1:length(regressor_var_names)
    if unconditional_model_dummy
        data_array_control=data_array_for_regression_stacked_by_variable(:,:,pos.([regressor_var_names{var_iter},'_lag_1']):pos.([regressor_var_names{var_iter},'_lag_',num2str(lag_number)]));
        data_array_for_regression=cat(3,data_array_for_regression,data_array_control);
    else
        %F(z)*X_{t-k}
        data_array_state_1 = data_array_for_regression_stacked_by_variable(:,:,pos.([regressor_var_names{var_iter},'_lag_1']):pos.([regressor_var_names{var_iter},'_lag_',num2str(lag_number)])).*...
            (repmat(data_array_for_regression_stacked_by_variable(:,:,Indicator_position),[1,1,lag_number]));
        %(1-F(z))*X_{t-k}
        data_array_state_2 = data_array_for_regression_stacked_by_variable(:,:,pos.([regressor_var_names{var_iter},'_lag_1']):pos.([regressor_var_names{var_iter},'_lag_',num2str(lag_number)])).*...
            (repmat(1-data_array_for_regression_stacked_by_variable(:,:,Indicator_position),[1,1,lag_number]));
        
        data_array_for_regression=cat(3,data_array_for_regression,data_array_state_1,data_array_state_2);
    end
end

%% Run regression

if unconditional_model_dummy == 1
    theta_mat=NaN(1+max_irf_horizon,1,length(dependent_var_names));
    se_mat=NaN(1+max_irf_horizon,1,length(dependent_var_names));
else
    theta_mat=NaN(1+max_irf_horizon,2,length(dependent_var_names));
    se_mat = NaN(1+max_irf_horizon,2,length(dependent_var_names));
    test_stat_mat = NaN(1+max_irf_horizon,length(dependent_var_names));
    test_p_mat = NaN(1+max_irf_horizon,length(dependent_var_names));
end

for var_iter=1:length(dependent_var_names)
    for horizon_iter = 0:max_irf_horizon
        %construct matrices without NaN
        if horizon_iter==0
            dependent_variable=squeeze(data_array_for_regression_stacked_by_variable(:,:,pos.(dependent_var_names{var_iter})));
            yhat_full=NaN(size(dependent_variable));
            [dependent_variable_for_run,data_array_for_regression_for_run,time_indices_non_NaN,country_indices_non_NaN]=create_regression_matrices_no_NaN(dependent_variable,data_array_for_regression,data_array_for_regression_stacked_by_variable,pos,country_indicator_names_mapping,time_fixed_effects_dummy,country_fixed_effects_dummy);
            % if var_iter==3 && horizon_iter==0
            %    generate_descriptives_table;
            % end
            if var_iter==3 && horizon_iter==0
               generate_nobs_ncountry;
            end
            %run regression
            [theta,stdDK,~,CovDK,yhat] = HszDk5cPs(dependent_variable_for_run,ones(size(dependent_variable_for_run,1),1),data_array_for_regression_for_run,1,7,1);
            
            if add_residual_dummy && ~strcmp(impulse_name,dependent_var_names(var_iter)) %when regression G on itself, residuals are 0
                yhat_full(time_indices_non_NaN,country_indices_non_NaN)=yhat;
                resids=dependent_variable-yhat_full;
            else
                resids=[];
            end
        else
            dependent_variable=squeeze(data_array_for_regression_stacked_by_variable(:,:,pos.([dependent_var_names{var_iter},'_lead_',num2str(horizon_iter)])));
            yhat_full=NaN(size(dependent_variable));
            
            [dependent_variable_for_run,data_array_for_regression_for_run,time_indices_non_NaN,country_indices_non_NaN]=create_regression_matrices_no_NaN(dependent_variable,cat(3,data_array_for_regression,resids),data_array_for_regression_stacked_by_variable,pos,country_indicator_names_mapping,time_fixed_effects_dummy,country_fixed_effects_dummy);
            %run regression
            [theta,stdDK,~,CovDK,yhat] = HszDk5cPs(dependent_variable_for_run,ones(size(dependent_variable_for_run,1),1),data_array_for_regression_for_run,1,7,1);
            if add_residual_dummy
                yhat_full(time_indices_non_NaN,country_indices_non_NaN)=yhat;
                resids=dependent_variable-yhat_full;
            else
                resids=[];
            end
        end
        
        %save coefficients
        if unconditional_model_dummy
            theta_mat(1+horizon_iter,:,var_iter) = theta(1,1);
            se_mat(1+horizon_iter,:,var_iter)  = stdDK(1,1)';
        else
            theta_mat(1+horizon_iter,1:2,var_iter) = theta(1:2,1)';
            se_mat(1+horizon_iter,:,var_iter)  = stdDK(1:2,1)';
            Rmat = zeros(1,length(theta));
            Rmat(1,1:2) = [1,-1];
            Wstat = (Rmat*theta)'/(Rmat*CovDK*Rmat')*(Rmat*theta);
            test_stat_mat(1+horizon_iter, var_iter) = Wstat;
            test_p_mat(1+horizon_iter, var_iter) = 1 - chi2cdf(Wstat,1);
        end
    end
end

%% Plot IRFs
fontsize = 6
if length(dependent_var_names) == 3
    subplot_index = [1,2,3];
    startsp = 1;
elseif length(dependent_var_names) == 4
    subplot_index = [2,4,1,5];
    startsp = 1;
else
    error('Specify Subplot Layout')
end


if sign_dummy == -1
    scaling = -0.25; % cut in government spending
    sign_plottitle = 'Gov. consumption cut';
elseif sign_dummy == 1
    scaling = 0.25; % cut in government spending
    sign_plottitle = 'Gov. consumption hike';
elseif sign_dummy == 0
    scaling = -0.25; % cut in government spending
    sign_plottitle = 'Gov. consumption cut (symmetric model)';
end

nptsvar=max_irf_horizon;
signific = 0.9;
confidence  = norminv(signific+(1-signific)/2,0,1);
mycolors = [1 1 1; .85 .85 .85];



main_fig=figure('Name',[Figure_name_split,': ',Figure_name_indicator]);
colormap(mycolors);

for sp = startsp:length(dependent_var_names)
    figure(main_fig)
    
    subplot(3,3,subplot_index(sp))
    
    if ~unconditional_model_dummy
        ci1a = scaling*(theta_mat(:,1,sp)+confidence*se_mat(:,1,sp));
        ci2a = scaling*(theta_mat(:,1,sp)-confidence*se_mat(:,1,sp));
        ci1b = scaling*(theta_mat(:,2,sp)+confidence*se_mat(:,2,sp));
        ci2b = scaling*(theta_mat(:,2,sp)-confidence*se_mat(:,2,sp));
    else
        ci1a=0;
        ci2a=0;
        ci1b = scaling*(theta_mat(:,1,sp)+confidence*se_mat(:,1,sp));
        ci2b = scaling*(theta_mat(:,1,sp)-confidence*se_mat(:,1,sp));
    end
    
    topa = max(ci1a,ci2a);
    bottoma = min(ci1a,ci2a);
    topb = max(ci1b,ci2b);
    bottomb = min(ci1b,ci2b);
    
    area(0:nptsvar,[bottomb, topb-bottomb],'EdgeColor','none','ShowBaseline','off','FaceColor','flat');
    hold on
    if ~unconditional_model_dummy
        fn=plot(0:nptsvar,scaling*theta_mat(:,2,sp),'LineWidth', 1);
        hold on
        fs=plot(0:nptsvar,scaling*theta_mat(:,1,sp),'r--','LineWidth', 1);
        hold on
        plot(0:nptsvar,bottoma','r:',0:nptsvar,topa','r:','LineWidth', 1)
    else
        fs=plot(0:nptsvar,scaling*theta_mat(:,1,sp),'LineWidth', 1);
    end
    hline(0,'k:')
    xlim([0 max_irf_horizon])
    box on;set(gca,'xTick',0:max_irf_horizon,'Layer','top','FontSize',fontsize);
    if sp == 1
        if ~unconditional_model_dummy
            legend([fn fs],State_name_vector,'Location','South')
        end
    elseif sp == 2 && length(dependent_var_names)==4
        %             if ~unconditional_model_dummy
        %                 legend([fn fs],State_name_vector,'Location','South')
        %             end
    elseif sp == 3 && strcmp(impulse_name,'OECD_Forecast_Error_g')
        legend([fn fs],State_name_vector,'Location','SouthEast')
    end
    title(dependent_var_plottitles(sp),'FontSize',fontsize)
    ylabel(dependent_var_ylabels(sp),'FontSize',fontsize)
    xlabel('quarters','FontSize',fontsize)
end
